"""
@Author: ISSCA_NEW
@Time:  16:43
@User: DELL
@File: attention.py
@Software: PyCharm
@Attention: 安耐毁誉，八风不动
            利、衰、毁、誉、称、讥、苦、乐

粗制滥造的拙作
"""
from Tools import base_tools as bt
from Tools import mathtools as mt
import netCDF4 as nc
import datetime
import numpy as np
import os
from itertools import chain
import time as ti

class write_nc:
    """
    将匹配好的数据保存到NC文件中
    """
    def __init__(self, filename: str):
        self.ncfile = nc.Dataset(filename, "w")

    def write_one_list(self, x: list, name: str,data_type:type or str):
        x_dim: nc.Dimension = self.ncfile.createDimension(name, len(x))
        x_var: nc.Variable = self.ncfile.createVariable(name, data_type, (name))
        x_var[:] = x
    def close(self):
        self.ncfile.close()


class cy_data:
    def __init__(self):
        self.lat: list = []#L2风速的纬度
        self.lon: list = []#L2风速的精度
        self.index: list = []
        self.channel: list = []#CYGNSS卫星的DDMI接收通道
        self.az: list = []#DDM方位角
        self.inc: list = []#信号入射角
        self.mss: list = []
        self.ws: list = []#yslf_wind_speed 详细见CYGNSS数据说明
        self.swh: list = []#有效波高，注意20210514及以前的swh来源于ERA-5，20210515后的swh来源于NOAA
        self.rcg: list = []#距离校正增益
        self.wsaz: list = []#ERA-5风场数据中的风向
        self.ews: list = []#ERA-5提供的风速
        self.cynum: list = []#接收反射信号的CYGNSS卫星号
        self.time: list = []#L2V3.1数据中的时间
        self.skt: list = []#ERA-5数据中的海表温度，顺手匹配的，本来还想匹配降雨、海冰什么的，但那些时空分辨率太低
        self.dis:list = []#L2V3.1风速点经纬度距离ERA-5风速点的空间距离
        self.tdis:list=[]#L2V3.1风速点经纬度距离ERA-5风速点的时间距离
        self.iws:list = []#ERA-5中的瞬时风速
        self.assort: dict = {
            "1": {
                "time": [], "lat": [], "lon": [], "az": [], "inc": [], "mss": [],
                "ws": [], "swh": [], "rcg": [], "cynum": [], "channel": [],
                "ews": [], "wsaz": [], "mwd": [], "skt": [], "crr": [],"dis":[],"tdis":[],"iws":[],
            },
            "2": {
                "time": [], "lat": [], "lon": [], "az": [], "inc": [], "mss": [],
                "ws": [], "swh": [], "rcg": [], "cynum": [], "channel": [],
                "ews": [], "wsaz": [], "mwd": [], "skt": [], "crr": [],"dis":[],"tdis":[],"iws":[],
            },
            "3": {
                "time": [], "lat": [], "lon": [], "az": [], "inc": [], "mss": [],
                "ws": [], "swh": [], "rcg": [], "cynum": [], "channel": [],
                "ews": [], "wsaz": [], "mwd": [], "skt": [], "crr": [],"dis":[],"tdis":[],"iws":[],
            },
            "4": {
                "time": [], "lat": [], "lon": [], "az": [], "inc": [], "mss": [],
                "ws": [], "swh": [], "rcg": [], "cynum": [], "channel": [],
                "ews": [], "wsaz": [], "mwd": [], "skt": [], "crr": [],"dis":[],"tdis":[],"iws":[],
            },
            "5": {
                "time": [], "lat": [], "lon": [], "az": [], "inc": [], "mss": [],
                "ws": [], "swh": [], "rcg": [], "cynum": [], "channel": [],
                "ews": [], "wsaz": [], "mwd": [], "skt": [], "crr": [],"dis":[],"tdis":[],"iws":[],
            },
            "6": {
                "time": [], "lat": [], "lon": [], "az": [], "inc": [], "mss": [],
                "ws": [], "swh": [], "rcg": [], "cynum": [], "channel": [],
                "ews": [], "wsaz": [], "mwd": [], "skt": [], "crr": [],"dis":[],"tdis":[],"iws":[],
            },
            "7": {
                "time": [], "lat": [], "lon": [], "az": [], "inc": [], "mss": [],
                "ws": [], "swh": [], "rcg": [], "cynum": [], "channel": [],
                "ews": [], "wsaz": [], "mwd": [], "skt": [], "crr": [],"dis":[],"tdis":[],"iws":[],
            },
            "8": {
                "time": [], "lat": [], "lon": [], "az": [], "inc": [], "mss": [],
                "ws": [], "swh": [], "rcg": [], "cynum": [], "channel": [],
                "ews": [], "wsaz": [], "mwd": [], "skt": [], "crr": [],"dis":[],"tdis":[],"iws":[],
            },
        }


    def sort(self):
        """
        由于CYGNSS的L2V3.1中,数据是按照时间排列的，后续进行L2和L1匹配时，需要同时加载8颗卫星的L1数据，十分占用内存。
        所有，我们将L2与ERA-5匹配后的数据，按照cy卫星号进行分类，这样在后续与L1匹配时，可以只导入1颗卫星的L1文件，匹配后导入下一颗卫星的L1数据
        """
        self.cynum = np.array(self.cynum)
        s_local = [
            [i for i, v in enumerate(np.where(self.cynum == 1, 1, 0).tolist()) if v],
            [i for i, v in enumerate(np.where(self.cynum == 2, 1, 0).tolist()) if v],
            [i for i, v in enumerate(np.where(self.cynum == 3, 1, 0).tolist()) if v],
            [i for i, v in enumerate(np.where(self.cynum == 4, 1, 0).tolist()) if v],
            [i for i, v in enumerate(np.where(self.cynum == 5, 1, 0).tolist()) if v],
            [i for i, v in enumerate(np.where(self.cynum == 6, 1, 0).tolist()) if v],
            [i for i, v in enumerate(np.where(self.cynum == 7, 1, 0).tolist()) if v],
            [i for i, v in enumerate(np.where(self.cynum == 8, 1, 0).tolist()) if v],
        ]
        self.num_local :list =list(chain(*[[i+1 for _ in range(0,len(s_local[i]))] for i in range(0,len(s_local))]))
        for i in range(0, len(s_local)):
            for j in range(0, len(s_local[i])):
                self.assort[str(i + 1)]["time"].append(self.time[s_local[i][j]])
                self.assort[str(i + 1)]["lat"].append(self.lat[s_local[i][j]])
                self.assort[str(i + 1)]["lon"].append(self.lon[s_local[i][j]])
                self.assort[str(i + 1)]["az"].append(self.az[s_local[i][j]])
                self.assort[str(i + 1)]["inc"].append(self.inc[s_local[i][j]])
                self.assort[str(i + 1)]["mss"].append(self.mss[s_local[i][j]])
                self.assort[str(i + 1)]["ws"].append(self.ws[s_local[i][j]])
                self.assort[str(i + 1)]["swh"].append(self.swh[s_local[i][j]])
                self.assort[str(i + 1)]["rcg"].append(self.rcg[s_local[i][j]])
                self.assort[str(i + 1)]["cynum"].append(self.cynum[s_local[i][j]])
                self.assort[str(i + 1)]["channel"].append(self.channel[s_local[i][j]])
                self.assort[str(i + 1)]["ews"].append(self.ews[s_local[i][j]])
                self.assort[str(i + 1)]["wsaz"].append(self.wsaz[s_local[i][j]])
                self.assort[str(i + 1)]["skt"].append(self.skt[s_local[i][j]])
                self.assort[str(i + 1)]["dis"].append(self.dis[s_local[i][j]])
                self.assort[str(i + 1)]["tdis"].append(self.tdis[s_local[i][j]])
                self.assort[str(i + 1)]["iws"].append(self.iws[s_local[i][j]])

    def save_nc(self, out_file_name):
        """
        将匹配好的数据保存到NC文件中
        """
        self.snc = write_nc(out_file_name)
        self.snc.write_one_list(self.num_local,"cynum","u1")
        self.snc.write_one_list(self.time,"time","f")
        self.snc.write_one_list(self.lat,"lat","f")
        self.snc.write_one_list(self.lon,"lon","f")
        self.snc.write_one_list(self.az,"az","f")
        self.snc.write_one_list(self.inc, "inc", "f")
        self.snc.write_one_list(self.mss, "mss", "f")
        self.snc.write_one_list(self.ws, "ws", "f")
        self.snc.write_one_list(self.swh, "swh", "f")
        self.snc.write_one_list(self.rcg, "rcg", "f")
        self.snc.write_one_list(self.channel, "channel", "f")
        self.snc.write_one_list(self.ews,"ews","f")
        self.snc.write_one_list(self.wsaz, "wsaz", "f")
        self.snc.write_one_list(self.skt, "skt", "f")
        self.snc.write_one_list(self.dis,"dis","f")
        self.snc.write_one_list(self.tdis, "tdis", "f")
        self.snc.write_one_list(self.iws,"iws","f")

        self.snc.close()


def get_time(delta_time: int) -> datetime.datetime:
    """
    ERA-5数据的时间为现在时间距离1900, 1, 1, 0, 0, 0的小时计数，将计数时间转换成UTC时间
    不过我们好像并未使用这个函数，在实际操作中，我们将L2数据按月分类到文件夹中，ERA-5数据按月下载，这样只需把ERA-5数据中的时间，
    按照列表位置/24h，即可获得月内日期
    """
    initial_time: datetime.datetime = datetime.datetime(1900, 1, 1, 0, 0, 0)
    return initial_time + datetime.timedelta(hours=delta_time)



def read_cy_nc(file_path: str) -> list:
    """
    读取cy的nc文件，L2V3.1
    """
    content = nc.Dataset(file_path)
    time = content['sample_time'][:].tolist()
    lat = content['lat']
    lon = content['lon']
    channel = content['ddm_channel']
    az = content['azimuth_angle']
    inc = content['incidence_angle']
    mss = content['mean_square_slope']
    ws = content['yslf_wind_speed']
    rcg = content['range_corr_gain']
    cynum = content['spacecraft_num']
    swh = content['swh']

    """
    tr为时间断点，具体值为[0，0.5,1.5,2.5,...,23.5]
    由于ERA-5的时间分辨率为1h，这样的话将匹配的时间标准设置为0.5h，由tr判断cy风速点属于ERA-5哪个时间段，简单粗暴
    另外，记录了时间距离，以后可以根据时间距离设置更严格的时间匹配标准，无需再与ERA-5数据进行匹配
    """
    tr = [(i + 0.5) * 3600 for i in range(0, 24)]
    tr.insert(0, 0)
    
    """
    根据tr,找到L2时间序列中对应时间段的位置，如UTC的(0.5h,1.5h)的位置，这样就可以直接切片获取数据

    代码中的binary_search函数为二分法
    """

    tr_value = [(bt.binary_search(time, tr[i]), bt.binary_search(time, tr[i + 1]))
                for i in range(0, len(tr) - 1)]
    """
    根据tr_value的时间分割，创建空列表，名为son，用来存放时间段的数据
    """
    son = [[] for _ in range(0, len(tr_value))]
    for i in range(0, len(son)):
        son[i].append({
                        "time":time[tr_value[i][0]:tr_value[i][1]],
                       "lat": lat[tr_value[i][0]:tr_value[i][1]],
                       "lon": lon[tr_value[i][0]:tr_value[i][1]],
                       "channel": channel[tr_value[i][0]:tr_value[i][1]],
                       "az": az[tr_value[i][0]:tr_value[i][1]],
                       "inc": inc[tr_value[i][0]:tr_value[i][1]],
                       "mss": mss[tr_value[i][0]:tr_value[i][1]],
                       "ws": ws[tr_value[i][0]:tr_value[i][1]],
                       "rcg": rcg[tr_value[i][0]:tr_value[i][1]],
                        "cynum":cynum[tr_value[i][0]:tr_value[i][1]],
                        "swh":swh[tr_value[i][0]:tr_value[i][1]],}
                       )
    return son


def get_time_dis(value:float)->float:
    """
    计算时间距离
    """
    mid = (value/3600)%1
    if mid>=0.5:
        return 1-mid
    else:
        return mid


"""
此处的timeshow功能为：获取函数运行时间。具体代码为：
def timeshow(func):
    from time import time
    def newfunc(*arg, **kw):
        t1 = time()
        res = func(*arg, **kw)
        t2 = time()
        print(f"{func.__name__: >10} : {t2 - t1:.6f} sec")
        return res
    return newfunc
"""
@bt.timeshow
def read_ec(file_path: str, cy_file_path: str,out_put_file):
    """
    加载ERA-5数据，此处的ERA-5数据为一个月30（31、28、29）天的数据。
    注意，我们在下载ERA-5数据时，为了节省内存，下载时进行了区域限制，即南北纬40°区域内的数据
    """
    content: nc.Dataset = nc.Dataset(file_path)
    time: list = content['time'][:].tolist()
    """
    将ERA-5数据中的时间，分割成每日时间标记
    """
    moon_day: list = [(i * 24, (i + 1) * 24) for i in range(0, int(len(time) / 24))]
    """
    这里lat的列表，我们忘记了为什么要用这么愚蠢的写法，好像是改某个bug时
    将lat倒序是为了方便使用二分法
    """
    lat: list = [i for i in content['latitude'][::-1]]
    lon :list = content['longitude'][:].tolist()
    """
    uws为u方向的风速分量，vws同理
    skt为海水表面温度
    iws为海面瞬时最大风速
    """
    uws: np.ma.core.MaskedArray = content['u10'][:]
    vws: np.ma.core.MaskedArray = content['v10'][:]
    skt: np.ma.core.MaskedArray = content['skt'][:]
    iws :np.ma.core.MaskedArray = content['i10fg'][:]
    del content

    """
    cy_file_path为存放L2数据的文件夹，一个文件夹里为1个月的数据
    """
    cylist: list = os.listdir(cy_file_path)

    for cy_file in cylist:
        a = ti.time()
        son = read_cy_nc(cy_file_path + cy_file)#读取L2数据
        """
        此处的cy_file[16:18]为L2数据文件名中的日期，转换成int并-1，-1的原因见moon_day列表
        """
        day_select: int = int(cy_file[16:18]) - 1
        """获取对应日期的ERA-5数据"""
        miduws: list = uws[moon_day[day_select][0]:moon_day[day_select][1]].tolist()
        midvws: list = vws[moon_day[day_select][0]:moon_day[day_select][1]].tolist()
        midskt: list = skt[moon_day[day_select][0]:moon_day[day_select][1]].tolist()
        midiws :list = iws[moon_day[day_select][0]:moon_day[day_select][1]].tolist()

        """新建cy_data类，用来存储匹配结果数据"""
        cyclass = cy_data()

        for i in range(0, 24):
        """
        以24h为循环
        """
            cylat: list = son[i][0]['lat'].tolist()
            cylon: list = son[i][0]['lon'].tolist()
            for j in range(0, len(cylat)):
                try:
                    if  abs(cylat[j]) < 40 :#去除南北纬40°以外的点
                        """
                        CYGNSS数据中，lon范围为[0,360]
                        ERA-5中，lon范围为[-180,180]
                        所有为了统一，将CY的lon转换成ERA-5格式的lon，命名为r_lon
                        一个经度，竟然有这么多中写法！
                        """
                        r_lon: float = cylon[j] if cylon[j] <= 180 else cylon[j] - 360

                        """
                        divide_local_lon为cy风速点经度在ERA-5经度列表中的位置，
                        +180的是因为ERA-5的经度列表为[-180,-177.5,....,180],*4的原因是ERA-5的经度分辨率为0.25°
                        """
                        divide_local_lon: int = int(round(r_lon / 0.25, 0)) + 180 * 4
                        divide_local_lat: int = int(round((cylat[j] + 40) / 0.25, 0))
                        elat = lat[len(lat) - divide_local_lat]
                        elon = lon[divide_local_lon]

                        
                        d_uws = miduws[i][len(lat) - divide_local_lat-1][divide_local_lon]
                        d_vws = midvws[i][len(lat) - divide_local_lat-1][divide_local_lon]
                        d_skt = midskt[i][len(lat) - divide_local_lat-1][divide_local_lon]
                        d_iws = midiws[i][len(lat) - divide_local_lat-1][divide_local_lon]
                        res_elat = lat[::-1][len(lat) - divide_local_lat-1]
                        res_elon = lon[divide_local_lon]
                        

                        if  son[i][0]['mss'][j] !="--" and son[i][0]['ws'][j] !="--" :
                                cyclass.wsaz.append(mt.az(0,0,d_uws,d_vws))
                                cyclass.lat.append(cylat[j])
                                cyclass.lon.append(cylon[j])
                                cyclass.time.append(son[i][0]['time'][j])
                                cyclass.channel.append(son[i][0]['channel'][j][0])
                                cyclass.az.append(son[i][0]['az'][j])
                                cyclass.inc.append(son[i][0]['inc'][j])
                                cyclass.mss.append(son[i][0]['mss'][j])
                                cyclass.ws.append(son[i][0]['ws'][j])
                                cyclass.rcg.append(son[i][0]['rcg'][j])
                                cyclass.ews.append(np.sqrt(d_uws ** 2 + d_vws ** 2))
                                cyclass.cynum.append(son[i][0]['cynum'][j])
                                cyclass.swh.append(son[i][0]['swh'][j])
                                cyclass.skt.append(d_skt)
                                cyclass.tdis.append(get_time_dis(son[i][0]['time'][j]))
                                cyclass.dis.append(mt.distance((res_elat,res_elon),(cylat[j],cylon[j])))
                                cyclass.iws.append(d_iws)
                except:
                    pass
        cyclass.sort()
        cyclass.save_nc(out_put_file+cy_file)
        b = ti.time()
        print((b-a)/60,len(cyclass.lat),cy_file)



read_ec(file_path= "G:/ERA5/2023/202309.nc", cy_file_path="G:/CYdata/L2V3.1/2023/09/",out_put_file="G:/MateL2/2023/09/")
read_ec(file_path= "G:/ERA5/2023/202310.nc", cy_file_path="G:/CYdata/L2V3.1/2023/10/",out_put_file="G:/MateL2/2023/10/")
# read_ec(file_path= "G:/ERA5/2023/202306.nc", cy_file_path="G:/CYdata/L2V3.1/2023/06/",out_put_file="G:/MateL2/2023/06/")
# read_ec(file_path= "G:/ERA5/2023/202307.nc", cy_file_path="G:/CYdata/L2V3.1/2023/07/",out_put_file="G:/MateL2/2023/07/")